import grpc
import prime_pb2
import prime_pb2_grpc
from concurrent.futures import ThreadPoolExecutor

# 定义工作节点列表
host_list = ['82.156.163.134:50051']
#host_list = ['localhost:50051']
#可以请求不同服务器多个端口、也可以请求同一个服务器多个端口，只要服务器开启了多个端口，就可以请求
#请求同一个服务器多个端口，可以实现负载均衡

def split_range(start, end, num_intervals):
    interval_size = (end - start) // num_intervals
    tasks = []
    for i in range(num_intervals):
        interval_start = start + i * interval_size
        interval_end = interval_start + interval_size
        if i == num_intervals - 1:
            interval_end = end
        tasks.append((interval_start, interval_end))
    return tasks

def main():
    # 为每个工作节点创建存根(stub)
    stubs = [
        prime_pb2_grpc.PrimeServiceStub(grpc.insecure_channel(host))
        for host in host_list
    ]

    tasks = split_range(10000, 50000, 5)

    with ThreadPoolExecutor() as executor:
        # 轮询分配任务到不同节点
        futures = [
            executor.submit(
                stubs[i % len(stubs)].CalculatePrimes,  # 轮询选择存根
                prime_pb2.TaskRequest(start=s, end=e)
            )
            for i, (s, e) in enumerate(tasks)
        ]

        primes = []
        for future in futures:
            response = future.result()
            primes.extend(response.primes)

    print(f"Found {len(primes)} primes: {sorted(primes)}")

if __name__ == '__main__':
    main()